# -*- coding: utf-8 -*-
"""
Created on Thu Jun  4 16:52:51 2020

@author: Farman
"""

from osgeo import gdal
import numpy

'''
GeoTiff data is record rows first.
'''


def save_array_to_geotiff(tiff_file, data_array, 
                              geo_transform=None, 
                              meta_data=None, 
                              projection=None, 
                              color_table=None):
        '''
        '''
        type_dic = {
            'int8':gdal.GDT_Byte,
            'int16':gdal.GDT_Int16,
            'int32':gdal.GDT_Int32,
            'uint8':gdal.GDT_Byte,
            'uint16':gdal.GDT_UInt16,
            'uint32':gdal.GDT_UInt32,
            'float32':gdal.GDT_Float32,
            'float64':gdal.GDT_Float64,
            'cint16':gdal.GDT_CInt16,   # complex numbers
            'cint32':gdal.GDT_CInt32,
            'complex32':gdal.GDT_CFloat32, # maybe numpy does not has this
            'complex64':gdal.GDT_CFloat64,
            'unknown':gdal.GDT_Unknown, # really unknown :)
        }
        
        if len(data_array.shape) == 3:
            layers, rows, columns = data_array.shape
        
        else:
            rows, columns = data_array.shape
            layers = 1 # for American ESACCI data only.
        
        dtype = type_dic.get(str(data_array.dtype))
        driver = gdal.GetDriverByName('Gtiff')
        outRaster = driver.Create(tiff_file, columns, rows, layers, dtype)
        #'CreateLayer' ignore, layers is initialized in driver.Create().
    
        #outRaster.SetDescription(ds.GetDescription())
        #'SetDescription', just file name by default, ignored.    
        #'CreateMaskBand', not use yet.    
        #'SetGCPs', empty cell in ds, ignored for all.    
        
        #'SetGeoTransform' !!! should be calculated each by each.
        if geo_transform is not None:
            outRaster.SetGeoTransform(geo_transform)
        
        #'SetMetadata',
        if meta_data is not None:
            outRaster.SetMetadata(meta_data)
        
        #'SetMetadataItem', not sure yet.
        
        #'SetProjection',
        if projection is not None:
            outRaster.SetProjection(projection)
        
        #'SetSpatialRef', unknown yet.    
        #'SetStyleTable', empty in ds, so ignored.    
        #'StartTransaction', unknown yet.
        #'TestCapability', useless?
        
        for band in range(1, 1 + layers):
            outband = outRaster.GetRasterBand(band)    
            outband.WriteArray(data_array[band - 1])
        
        if color_table is not None:
            outband.SetColorTable(color_table)
            
        return


class GeoTransform:
    def __init__(self, geo_transform):
        self.geo_transform = geo_transform
        m = list(geo_transform)
        self.a = numpy.mat([[m[1], m[2]], [m[4], m[5]]])
        self.b = numpy.mat([[m[0]], [m[3]]])
        return
    
    
    def get_projected(self, x, y):
        p = self.a * numpy.mat([[x], [y]]) + self.b
        p = numpy.array(p)[:,0]
        return p[0], p[1]
    
    
    def get_projected_mat(self, mat_xy):
        '''
        xy_mat = [[x,y],[x,y],[x,y],...]
        '''
        xy_mat = numpy.mat(mat_xy).getT()
        xy_mat = self.a * xy_mat + self.b
        xy_mat = numpy.array(xy_mat)
        return xy_mat.T
    
    
    def get_reverse_project(self, x, y):
        '''
        lon, lat etc. -> x, y in scene.
        '''
        xyp = numpy.mat([[x], [y]])
        xy  = self.a.getI() * (xyp - self.b)
        [x, y] = numpy.array(xy[:, 0])
        return int(x+0.5), int(y+0.5)
    
    
    def get_reverse_project_mat(self, mat_xy):
        '''
        [[lon, lat]] etc. -> [[x, y]] in scene.
        xy_mat = [[x,y],[x,y],[x,y],...]
        '''
        xy_mat = numpy.mat(mat_xy).getT()
        xy_mat = self.a.getI() * (xy_mat - self.b)
        xy_mat += numpy.mat([[0.5], [0.5]])
        xy_mat = numpy.array(xy_mat, dtype='int')
        return xy_mat.T
    
    def get_geo_transform_at_origin(self, origin_x, origin_y, resample_rate=1):
        '''
        resample_rate > 1 if want to split pixels.
        resample_rate < 1 if want to merge pixels.
        '''
        p = self.a * numpy.mat([[origin_x], [origin_y]]) + self.b
        p = numpy.array(p)[:,0]
        a = self.a / resample_rate
        
        return (p[0], a[0, 0], a[0, 1],
                p[1], a[1, 0], a[1, 1])
    
    
    
'''
geo_transform = (102.98371494013941,
 2.0720589871671455e-05,
 0.0,
 30.835714500118257,
 0.0,
 -2.0720589871671556e-05)

gt = GeoTransform(geo_transform)
x, y = 1000, 1000
x, y = gt.get_projected(x, y)
print(x, y)

x, y = gt.get_reverse_project(x, y)
print(x, y)

g = gt.get_geo_transform_at_origin(1000, 10000)
print(g)

g = gt.get_geo_transform_at_origin(1000, 10000, 10)
print(g)


xy_mat = numpy.array([[0,0], [100,100], [1000, 1000]])
print(xy_mat)

xy_mat = gt.get_projected_mat(xy_mat)
print(xy_mat)

xy_mat = gt.get_reverse_project_mat(xy_mat)
print(xy_mat)

#--------------------------------------------------------------------------'''

class GeoCoorTran:
    def __init__(self, geo_transform_from, geo_transform_to):
        self.geo_transformer_from = GeoTransform(geo_transform_from)
        self.geo_transformer_to   = GeoTransform(geo_transform_to)
        return
    
    #-------------------------------------------------------------------------
    
    def set_geo_transform(self, geo_transform_from, geo_transform_to):
        self.geo_transformer_from = GeoTransform(geo_transform_from)
        self.geo_transformer_to   = GeoTransform(geo_transform_to)
        return
    
    
    def set_geo_transform_from(self, geo_transform_from):
        self.geo_transformer_from = GeoTransform(geo_transform_from)
        return
    
    
    def set_geo_transform_to(self, geo_transform_to):
        self.geo_transformer_to   = GeoTransform(geo_transform_to)
        return
    
    #-------------------------------------------------------------------------
        
    def get_tran(self, x_from, y_from):
        xp, yp = self.geo_transformer_from.get_projected(x_from, y_from)
        return self.geo_transformer_to.get_reverse_project(xp, yp)

    
    def get_reverse_tran(self, x_to, y_to):
        xp, yp = self.geo_transformer_to.get_projected(x_to, y_to)
        return self.geo_transformer_from.get_reverse_project(xp, yp)
    
        
    def tran_mat(self, xy_mat):
        '''
        xy_mat = [[x,y],[x,y],[x,y],...]
        '''
        xy_mat = self.geo_transformer_from.get_projected_mat(xy_mat)
        return self.geo_transformer_to.get_reverse_project_mat(xy_mat)
    
    # END OF CLASS
    #-------------------------------------------------------------------------

    
        
'''
geo_transform_from = (102.98371494013941,
 2.0720589871671455e-05,
 0.0,
 30.835714500118257,
 0.0,
 -2.0720589871671556e-05)

geo_transform_to = (102.923907579582,
 2.0949831999999985e-05,
 0.0,
 30.983970110789,
 0.0,
 -2.0949832000000052e-05)

t = GeoCoorTran(geo_transform_from, geo_transform_to)

x, y = 1, 25
print(x, y)

x, y = t.get_tran(x, y)
print(x, y)

x, y = t.get_reverse_tran(x, y)
print(x, y)
#'''


class GeoTiffMan:
    def __init__(self, tiff_file=''):
        '''
        

        Parameters
        ----------
        tiff_file : TYPE, optional
            DESCRIPTION. The default is ''.

        Returns
        -------
        None.

        '''
        self.data_set = None
        self.__open_dataset__(tiff_file)
        
        self.type_dic = {
            'int8':gdal.GDT_Byte,
            'int16':gdal.GDT_Int16,
            'int32':gdal.GDT_Int32,
            'uint8':gdal.GDT_Byte,
            'uint16':gdal.GDT_UInt16,
            'uint32':gdal.GDT_UInt32,
            'float32':gdal.GDT_Float32,
            'float64':gdal.GDT_Float64,
            'cint16':gdal.GDT_CInt16,   # complex numbers
            'cint32':gdal.GDT_CInt32,
            'complex32':gdal.GDT_CFloat32, # maybe numpy does not has this
            'complex64':gdal.GDT_CFloat64,
            'unknown':gdal.GDT_Unknown, # really unknown :)
        }
        
        return
    
    
    def __open_dataset__(self, tiff_file):
        tiff_file = str(tiff_file)
        
        del self.data_set
        
        if tiff_file == '':
            self.dataset = None
        else:
            try:        
                self.data_set = gdal.Open(tiff_file)
                self.geo_transformer = GeoTransform(self.data_set.GetGeoTransform())
                self.meta_data = self.data_set.GetMetadata()
                self.projection = self.data_set.GetProjection()
                self.columns = self.data_set.RasterXSize
                self.rows    = self.data_set.RasterYSize

                # MAY NEED DEBUG HERE
                # correct for US ESACCI file.
                self.color_table = self.get_color_table(1) 
            except:
                self.data_set = None
                print('Failed to open %s'%tiff_file)
        return
    
    
    def read_bands(self, tiff_file=''):
        '''
        Read geotiff data.
        '''
        if tiff_file == '':
            pass
        else:
            self.__open_dataset__(tiff_file)
        
        if self.data_set == None:
            print('Dataset not ready.')
            return numpy.array([])
        
        bands = self.data_set.RasterCount
        width  = self.data_set.RasterXSize
        height = self.data_set.RasterYSize
        bands_data = []
        
        for band in range(1, bands + 1):
                band_data = self.data_set.GetRasterBand(band)
                bands_data.append(
                        band_data.ReadAsArray(0, 0, width, height))
        
        return numpy.array(bands_data)
    

    def read_band(self, band, tiff_file=''):
        '''
        Read geotiff data.
        '''
        if tiff_file == '':
            pass
        else:
            self.__open_dataset__(tiff_file)
        
        if self.data_set == None:
            print('Dataset not ready.')
            return numpy.array([])
        
        width  = self.data_set.RasterXSize
        height = self.data_set.RasterYSize
        band_data = self.data_set.GetRasterBand(band)
        return band_data.ReadAsArray(0, 0, width, height)


    def get_attrib(self, atrrib_name, tiff_file=''):
        if tiff_file == '':
            pass
        else:
            self.__open_dataset__(tiff_file)
        
        if self.data_set == None:
            print('Dataset not ready.')
            return 'Dataset not ready.'
    
        ds = self.data_set
    
        try:
            result = eval("ds.%s()"%atrrib_name)
        except:
            result = eval("ds.%s"%atrrib_name)
    
        return result
    
    
    def get_geo_transform(self, tiff_file=''):
        if tiff_file == '':
            pass
        else:
            self.__open_dataset__(tiff_file)
        
        if self.data_set == None:
            print('Dataset not ready.')
            return None
        
        return self.data_set.GetGeoTransform()
    
    
    def get_meta_data(self, tiff_file=''):
        if tiff_file == '':
            pass
        else:
            self.__open_dataset__(tiff_file)
        
        if self.data_set == None:
            print('Dataset not ready.')
            return None
        
        return self.data_set.GetMetadata()
    
    
    def get_projection(self, tiff_file=''):
        if tiff_file == '':
            pass
        else:
            self.__open_dataset__(tiff_file)
        
        if self.data_set == None:
            print('Dataset not ready.')
            return None
        
        return self.data_set.GetProjection()
    
    
    def get_color_table(self, band, tiff_file=''):
        if tiff_file == '':
            pass
        else:
            self.__open_dataset__(tiff_file)
        
        if self.data_set == None:
            print('Dataset not ready.')
            return 'Dataset not ready.'
        
        band_data = self.data_set.GetRasterBand(band)
        return band_data.GetColorTable()
    

    def get_columns(self):
        return self.columns


    def get_rows(self):
        return self.rows


    def get_size(self):
        '''
        return size_x, size_y
        '''
        return self.columns, self.rows
    
    
    def show_info(self, tiff_file=''):
        if tiff_file == '':
            pass
        else:
            self.__open_dataset__(tiff_file)
        
        if self.data_set == None:
            print('Dataset not ready.')
            return
        
        items = dir(self.data_set)
    
        for item in items:
            if item[0] != '_':
                print('\n', item, ':', end=' ')
    
                if item.find('Read') > -1:
                    print("Maybe huge data, ignored.")
                    continue
                
                try:
                    print(eval("self.data_set.%s()"%item))
                except:
                    print(eval("self.data_set.%s"%item))
        
        return


    def save_info(self, info_file, tiff_file=''):
        if tiff_file == '':
            pass
        else:
            self.__open_dataset__(tiff_file)
        
        if self.data_set == None:
            print('Dataset not ready.')
            return
        
        ds = self.data_set
        items = dir(ds)
        lines = []
    
        for item in items:
            if item[0] != '_':
                lines.append('\n\n%s :'%item)
                #print('\n\n%s :'%item)
                lines.append('-' * 30)
                
                if item.find('Read') > -1:
                    lines.append('\tMay be huge data. Skipped.')
                    continue
    
                try:
                    #print('%s'%eval("ds.%s()"%item))
                    lines.append('%s'%eval("ds.%s()"%item))
                except:
                    lines.append('%s'%eval("ds.%s"%item))
    
        lines = [line + '\n' for line in lines]
        
        with open(info_file, 'w') as f:
            f.writelines(lines)
            
        return


    def save_as_geotiff(self, dest_tiff_file, 
                        offset_x, 
                        offset_y,
                        size_x,
                        size_y,
                        source_tiff_file=''):
        '''
        '''
        source_tiff_file = str(source_tiff_file)
        dest_tiff_file   = str(dest_tiff_file)
        
        if source_tiff_file == '':
            pass
        else:
            self.__open_dataset__(source_tiff_file)
        
        if self.data_set == None:
            print('Dataset not ready.')
            print('Check validation of %s'%source_tiff_file)
            return
        
        data_array = self.data_set.ReadAsArray(offset_x, offset_y, size_x, size_y)
        #data_array = self.data_set.ReadAsArray()
        #return data_array
    
        if len(data_array.shape) == 2:
            data_array = numpy.array([data_array])
            
        bands, rows, columns = data_array.shape     
        dtype = self.type_dic.get(str(data_array.dtype))
        driver = gdal.GetDriverByName('Gtiff')
        outRaster = driver.Create(dest_tiff_file, columns, rows, bands, dtype)
        #'CreateLayer' ignore, layers is initialized in driver.Create().
        
        #outRaster.SetDescription(ds.GetDescription())
        #'SetDescription', just file name by default, ignored.    
        #'CreateMaskBand', not use yet.    
        #'SetGCPs', empty cell in ds, ignored for all.    
        
        #'SetGeoTransform' !!! should be calculated each by each.
        outRaster.SetGeoTransform(
            self.geo_transformer.get_geo_transform_at_origin(offset_x, offset_y))
               
        #'SetMetadata',
        outRaster.SetMetadata(self.meta_data)
        
        #'SetMetadataItem', not sure yet.
        
        #'SetProjection',
        outRaster.SetProjection(self.projection)
        
        #'SetSpatialRef', unknown yet.    
        #'SetStyleTable', empty in ds, so ignored.    
        #'StartTransaction', unknown yet.
        #'TestCapability', useless?
        
        for band in range(1, 1 + bands):
            outband = outRaster.GetRasterBand(band)
            # set color table
            outband.SetColorTable(self.color_table)
            outband.WriteArray(data_array[band - 1])
        
        return
    
    
    def save_as_resampled_geotiff(self, dest_tiff_file, 
                        offset_x, 
                        offset_y,
                        size_x,
                        size_y,
                        resample_rate=1,
                        source_tiff_file=''):
        '''
        '''
        if source_tiff_file == '':
            pass
        else:
            self.__open_dataset__(source_tiff_file)
        
        if self.data_set == None:
            print('Dataset not ready.')
            print('Check validation of %s'%source_tiff_file)
            return
        
        #data_array = self.data_set.ReadAsArray(offset_x, offset_y, size_x, size_y)
        data_array = self.data_set.ReadAsArray(offset_x, offset_y, size_x, size_y, 
                                       None, size_x * resample_rate, 
                                       size_y * resample_rate)
    
        if len(data_array.shape) == 2:
            data_array = numpy.array([data_array])
            
        bands, rows, columns = data_array.shape     
        dtype = self.type_dic.get(str(data_array.dtype))
        driver = gdal.GetDriverByName('Gtiff')
        outRaster = driver.Create(dest_tiff_file, 
                                  columns * resample_rate, 
                                  rows * resample_rate, bands, dtype)
        #'CreateLayer' ignore, layers is initialized in driver.Create().
        
        #outRaster.SetDescription(ds.GetDescription())
        #'SetDescription', just file name by default, ignored.    
        #'CreateMaskBand', not use yet.    
        #'SetGCPs', empty cell in ds, ignored for all.    
        
        #'SetProjection'
        outRaster.SetProjection(self.data_set.GetProjection())
        
        #'SetGeoTransform' !!! should be calculated each by each.
        outRaster.SetGeoTransform(
            self.geo_transformer.get_geo_transform_at_origin(
                    offset_x, offset_y, resample_rate))
               
        #'SetMetadata',
        outRaster.SetMetadata(self.meta_data)
        
        #'SetMetadataItem', not sure yet.
                
        #'SetSpatialRef', unknown yet.    
        #'SetStyleTable', empty in ds, so ignored.    
        #'StartTransaction', unknown yet.
        #'TestCapability', useless?
        
        for band in range(1, 1 + bands):
            outband = outRaster.GetRasterBand(band)
            # set color table
            outband.SetColorTable(self.color_table)
            outband.WriteArray(data_array[band - 1])
        
        print(outRaster.GetProjection())
        return 

#'''
#=========================================================================================================================================
file_name = r'/data/yule/disk2/ESACCI/ESACCI-LC-L4-LCCS-Map-300m-P1Y-2010-v2.0.7.tif'
r = GeoTiffMan(file_name)
'''
#r.show_info()
r.save_as_geotiff('1.tif', 100, 100, 100, 100)
da = r.save_as_resampled_geotiff('2.tif', 100, 100, 100, 100, 10)
#print(r.geo_transformer.geo_transform)

print(r.get_geo_transform('1.tif'))
print(r.get_geo_transform('2.tif'))
#print(r.get_meta_data('1.tif'))
#print(r.get_meta_data('2.tif'))
print(r.get_projection('1.tif'))
print(r.get_projection('2.tif'))

r = GeoTiffMan('2.tif')
print(r.get_projection())
r.show_info()
#--------------------------------------------------------------------------'''

if __name__ == '__main__':
    pass
